package com.lft.tree10.segment_tree;

/**
 * 线段树
 */
public class SegmentTreeDemo {
	public static void main(String[] args) {
		Integer[] a = {3, 5, 6, 8, 4, 2, 1, 7, 9};
		SegmentTree<Integer> segmentTree =
				new SegmentTree<Integer>(a, new Merger<Integer>() {
					@Override
					public Integer merger(Integer a, Integer b) {
						return a + b;
					}
				});
		System.out.println(segmentTree.getSize());
		System.out.println(segmentTree.query(2, 3));
	}
}

class SegmentTree<E> {
	/**
	 * 存储原始数据的数组
	 */
	private E[] data;
	/**
	 * 存储线段树的数组
	 */
	private E[] tree;
	
	private Merger<E> merger;
	
	/**
	 * 构造器
	 * @param arr    原始数组
	 * @param merger 自己实现的合并方法
	 */
	public SegmentTree(E[] arr, Merger<E> merger) {
		this.merger = merger;
		this.data = (E[]) new Object[arr.length];
		for (int i = 0; i < arr.length; i++) {
			data[i] = arr[i];
		}
		// 数组存储线段树需要开辟出 4 倍的空间。
		tree = (E[]) new Object[4 * arr.length];
		// 构建线段树
		buildSegment(0, 0, arr.length - 1);
	}
	
	/**
	 * 创建线段树的方法
	 * 树具有天然的递归特性，所以在创建线段树的过程中我们可以很方便快速的使用递归创建线段树。
	 * 但是，需要注意的是由于我们使用了泛型，所以合并左右孩子的方法需要我们自己传入
	 * @param treeIndex 节点对应索引
	 * @param left      左边界
	 * @param right     右边界
	 */
	private void buildSegment(int treeIndex, int left, int right) {
		// 划分到叶子节点了，构建完成。
		if (left == right) {
			tree[treeIndex] = data[left];
			return;
		}
		int leftChildIndex = leftChildIndex(treeIndex);
		int rightChildIndex = rightChildIndex(treeIndex);
		/**
		 * 求出 mid ，将 left ~ right 的区间划分成 left ~ mid 和 mid + 1 ~ right 的区间。
		 * 对划分后的两个区间再创建线段树。
		 */
		int mid = left + (right - left) / 2;
		buildSegment(leftChildIndex, left, mid);
		buildSegment(rightChildIndex, mid + 1, right);
		// 根据业务需求合并两个左右子树的值到父节点。
		tree[treeIndex] = merger.merger(tree[leftChildIndex], tree[rightChildIndex]);
	}
	
	/**
	 * 查询线段树
	 * @param left  查询的左边界
	 * @param right 查询的右边界
	 * @return
	 */
	public E query(int left, int right) {
		if (left < 0 || right < 0 || left > right || left > data.length || right > data.length) {
			throw new IllegalArgumentException("参数非法");
		}
		return query(0, 0, data.length - 1, left, right);
	}
	
	/**
	 * 以 treeIndex 为索引的线段树 [leftIndex,rightIndex] 的范围里，搜索 [queryLeft,queryRight] 的值
	 * @param treeIndex  节点索引
	 * @param leftIndex  左边界索引
	 * @param rightIndex 右边界索引
	 * @param queryLeft  查询的左边界
	 * @param queryRight 查询的右边界
	 * @return
	 */
	private E query(int treeIndex, int leftIndex, int rightIndex, int queryLeft, int queryRight) {
		// 左边界索引和查询的左边界相等，右边界索引和查询的右边界相等。当前树节点索引就是要查找的。
		if (leftIndex == queryLeft && rightIndex == queryRight) {
			return tree[treeIndex];
		}
		int mid = leftIndex + (rightIndex - leftIndex) / 2;
		int leftTreeIndex = leftChildIndex(treeIndex);
		int rightTreeIndex = rightChildIndex(treeIndex);
		
		// 如果要查询的左边界大于 mid +1，就到右子树中递归查找。
		if (queryLeft > mid + 1) {
			return query(rightTreeIndex, mid + 1, rightIndex, queryLeft, queryRight);
		}
		// 如果要查询的右边界小于 mid，就到左子树中递归查找。
		else if (queryRight < mid) {
			return query(leftTreeIndex, leftIndex, mid, queryLeft, queryRight);
		}
		// 查询区间跨越了左右孩子。
		else {
			return merger
					.merger(query(leftTreeIndex, leftIndex, mid, queryLeft, mid),
							query(rightTreeIndex, mid + 1, rightIndex, mid + 1, queryRight));
		}
	}
	
	/**
	 * 更新线段树的某一节点的值
	 * @param index    要更新的节点的索引
	 * @param newValue 新值
	 */
	public void updateOne(int index, E newValue) {
		if (index < 0 || index >= data.length) {
			throw new IllegalArgumentException("Index Out Of Bounds");
		}
		data[index] = newValue;
		// 更新父节点的值
		updateOne(0, 0, data.length - 1, index, newValue);
	}
	
	/**
	 * 更新父节点的值。
	 * @param treeIndex  根节点索引
	 * @param leftIndex  左边界
	 * @param rightIndex 右边界
	 * @param index      要更新的索引
	 * @param newValue   新值
	 */
	private void updateOne(int treeIndex, int leftIndex, int rightIndex, int index, E newValue) {
		// 左边界索引和查询的左边界相等，右边界索引和查询的右边界相等。当前树节点索引就是要查找的。
		if (leftIndex == rightIndex) {
			tree[treeIndex] = newValue;
			return;
		}
		int mid = leftIndex + (rightIndex - leftIndex) / 2;
		int leftTreeIndex = leftChildIndex(treeIndex);
		int rightTreeIndex = rightChildIndex(treeIndex);
		// 待更新节点在右子树中
		if (index >= mid + 1) {
			updateOne(rightTreeIndex, mid + 1, rightIndex, index, newValue);
		}
		// 待更新节点在左子树中
		else {
			updateOne(leftTreeIndex, leftIndex, mid, index, newValue);
		}
		// 更新父节点的值
		tree[treeIndex] = merger.merger(tree[leftTreeIndex], tree[rightTreeIndex]);
	}
	
	/**
	 * 获取下标为 index 的数据
	 * @param index
	 * @return
	 */
	public E get(int index) {
		if (index < 0 || index >= data.length) {
			throw new IllegalArgumentException("Index Out Of Bounds");
		}
		return data[index];
	}
	
	/**
	 * 获取元素的数量
	 * @return
	 */
	public int getSize() {
		return data.length;
	}
	
	/**
	 * 根据当前节点的索引获取左子节点的索引
	 * @param index
	 * @return
	 */
	private int leftChildIndex(int index) {
		return index * 2 + 1;
	}
	
	/**
	 * 根据当前节点的索引获取右子节点的索引
	 * @param index
	 * @return
	 */
	private int rightChildIndex(int index) {
		return index * 2 + 2;
	}
	
}

class Node implements Merger<Integer> {
	private int val;
	private int leftIndex;
	private int rightIndex;
	private Node leftChild;
	private Node rightChild;
	
	public Node(int leftIndex, int rightIndex) {
		this.val = 0;
		this.leftIndex = leftIndex;
		this.rightIndex = rightIndex;
		this.leftChild = null;
		this.rightChild = null;
	}
	
	@Override
	public Integer merger(Integer a, Integer b) {
		return a + b;
	}
}

interface Merger<E> {
	E merger(E a, E b);
}